import json
import pandas as pd
from flask import jsonify
from statsmodels.tsa.arima.model import ARIMA


def perform_arima_prediction(p, d, q, steps, data_file='./data/test2.json', forecast_file='./data/test3.json'):
    # 加载数据
    with open(data_file, 'r') as file1:
        data = json.load(file1)

    # 创建 DataFrame
    df = pd.DataFrame(data)
    df.set_index(pd.Index([i * 0.1 for i in range(len(df))]), inplace=True)

    # 拟合 ARIMA 模型
    model = ARIMA(df['y'], order=(p, d, q))
    fitted_model = model.fit()

    # 进行预测
    forecast = fitted_model.forecast(steps=steps)

    print(f"数据预测 {forecast}")  # 确认数据写入

    # 创建预测数据的 JSON 对象
    last_x = df.index[-1] + 0.1
    forecast_data = [{'x': round(last_x + i * 0.1, 1), 'y': int(round(y))} for i, y in enumerate(forecast)]

    # 保存预测数据到新的 JSON 文件
    with open(forecast_file, 'w') as file2:
        json.dump(forecast_data, file2, indent=4)

    print(f"Data written to {forecast_file}")  # 确认数据写入

    print("ARIMA 预测完成，并已保存到新的 JSON 文件。")

def get_AIC(n = 4, data_file='./data/purchasing/test2.json'):
    p_values = range(0, n)
    results = []

    # 加载数据
    with open(data_file, 'r') as file1:
        data = json.load(file1)

    # 创建 DataFrame
    df = pd.DataFrame(data)
    df.set_index(pd.Index([i * 0.1 for i in range(len(df))]), inplace=True)

    for p in p_values:
        try:
            model = ARIMA(df['y'], order=(p, 1, 1))
            fitted_model = model.fit()
            aic = fitted_model.aic
            results.append({'p': p, 'AIC': aic})
        except:
            continue
    return results

def get_AICb(n = 4, data_file='./data/purchasing/test2.json'):
    q_values = range(0, n)
    results = []

    # 加载数据
    with open(data_file, 'r') as file1:
        data = json.load(file1)

    # 创建 DataFrame
    df = pd.DataFrame(data)
    df.set_index(pd.Index([i * 0.1 for i in range(len(df))]), inplace=True)
    for q in q_values:
        try:
            model = ARIMA(df['y'], order=(1, q, 1))
            fitted_model = model.fit()
            aic = fitted_model.aic
            results.append({'q': q, 'AIC': aic})
        except:
            continue
    return results
    

# # Example usage
# if __name__ == '__main__':
#     perform_arima_prediction(1, 1, 1, 10, data_file='data/test2.json', forecast_file='data/test3.json')